# -*- coding: utf-8 -*-

import os
import argparse
import pandas as pd

from train import load_model
from utils import *
from auto_correct import Speller

# v2.0 添加拼写纠错模块
spell = Speller('ug')


# 导入近义词表

class Query(object):
    def __init__(self, args, model):
        self.args = args
        self.table_df = self.get_table()
        self.model = model

    def get_table(self):
        table_df = pd.DataFrame()
        if not self.args.exists_word:
            sim_type = 'all'
        else:
            sim_type = 'exists'
        file = os.path.join(table_path, sim_type, table_save_file[self.args.lang])
        if os.path.exists(file):
            table_df = pd.read_csv(file, header=None)
        table_df.columns = ['words', 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        return table_df

    # 查询目标词 text 的前10个近义词，若词表中没有结果，则从模型中计算
    def get_similar(self, text):
        sim = []
        # if text in self.table_df['words'].values:
        #     sim = self.table_df[self.table_df['words'] == text].values.tolist()
        #     sim = sim[0]
        if not sim:
            sim.append(text)
            sim_words = self.model.wv.most_similar(text, topn=10)
            for s in sim_words:
                # v2.0 添加拼写纠错模块
                s0 = spell.autocorrect_sentence(s[0], mode='exist')
                print(s0)
                if s0 not in [si[0] for si in sim] and s0 != '':
                    s = (s0, s[1])
                    sim.append(s)
        for w in sim:
            print(w)
        return sim

    # 查询目标输入 ./result/input_ug.txt 列表中所有词的前10个近义词，结果保存在 ./result/result_ug.csv
    def get_predict(self):
        input_df = pd.read_csv(input_file[self.args.lang], header=None)
        input_df.columns = ['words']
        similar = []
        for text in input_df['words']:
            sim = self.get_similar(text)
            similar.append(sim)

        similar_df = pd.DataFrame(similar)
        # similar_df = pd.merge(input_df, similar_df, how='outer', left_index=True, right_index=True)

        # 保存查询结果
        similar_df.to_csv(result_save_file[self.args.lang], encoding='utf-8-sig', index=0, header=0)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    args = get_args(parser)
    model = load_model(args)  # 导入模型

    query = Query(args, model)
    query.get_similar('گاز ئۇچاق')
    # query.get_predict()
